from asgiref.sync import sync_to_async
from django.utils import asyncio
from rest_framework.decorators import parser_classes
from rest_framework.parsers import MultiPartParser, JSONParser
from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework import status
from drf_yasg.utils import swagger_auto_schema
from drf_yasg import openapi
import base64
import numpy as np
import cv2
from django.core.files.uploadedfile import InMemoryUploadedFile
from .utils import model


# @sync_to_async
def async_predict(img):
    return model.predict(img)


class Base64DetectionAPI(APIView):
    # 处理 JSON 格式的请求体
    parser_classes = [JSONParser]

    @swagger_auto_schema(
        operation_description="YOLOv8 目标检测 (Base64 输入)",
        request_body=openapi.Schema(
            type=openapi.TYPE_OBJECT,
            required=['image'],
            properties={
                'image': openapi.Schema(
                    type=openapi.TYPE_STRING,
                    description='Base64 编码的图像数据',
                    example='...'
                )
            }
        ),
        responses={
            200: openapi.Response(
                description="检测结果",
                schema=openapi.Schema(
                    type=openapi.TYPE_OBJECT,
                    properties={
                        'results': openapi.Schema(
                            type=openapi.TYPE_ARRAY,
                            items=openapi.Schema(
                                type=openapi.TYPE_OBJECT,
                                properties={
                                    'class': openapi.Schema(type=openapi.TYPE_STRING),
                                    'confidence': openapi.Schema(type=openapi.TYPE_NUMBER),
                                    'box': openapi.Schema(
                                        type=openapi.TYPE_ARRAY,
                                        items=openapi.Schema(type=openapi.TYPE_NUMBER)
                                    )
                                }
                            )
                        )
                    }
                )
            )
        }
    )
    def post(self, request):
        # def post(self, request):
        try:
            image_data = request.data.get('image').split(',')[1]
            nparr = np.frombuffer(base64.b64decode(image_data), np.uint8)
            img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

            # 使用全局模型进行推理
            # results = await async_predict(img)
            results = async_predict(img)

            # 处理结果
            output = []
            for result in results:
                for box in result.boxes:
                    output.append({
                        'class': model.names[int(box.cls)],
                        'confidence': float(box.conf),
                        'box': box.xywhn.tolist()[0]
                    })

            return Response({'results': output}, status=status.HTTP_200_OK)

        except Exception as e:
            return Response({'error': str(e)}, status=status.HTTP_400_BAD_REQUEST)


class FileDetectionAPI(APIView):
    # 处理 JSON 格式的请求体
    parser_classes = [MultiPartParser]

    @swagger_auto_schema(
        operation_description="YOLOv8 目标检测 (文件上传)",
        manual_parameters=[
            openapi.Parameter(
                'file',
                openapi.IN_FORM,
                type=openapi.TYPE_FILE,
                required=True,
                description='上传的图像文件'
            )
        ],
        responses={
            200: openapi.Response(
                description="检测结果",
                schema=openapi.Schema(
                    type=openapi.TYPE_OBJECT,
                    properties={
                        'results': openapi.Schema(
                            type=openapi.TYPE_ARRAY,
                            items=openapi.Schema(
                                type=openapi.TYPE_OBJECT,
                                properties={
                                    'class': openapi.Schema(type=openapi.TYPE_STRING),
                                    'confidence': openapi.Schema(type=openapi.TYPE_NUMBER),
                                    'box': openapi.Schema(
                                        type=openapi.TYPE_ARRAY,
                                        items=openapi.Schema(type=openapi.TYPE_NUMBER)
                                    )
                                }
                            )
                        )
                    }
                )
            )
        }
    )
    def post(self, request):
        try:
            uploaded_file: InMemoryUploadedFile = request.FILES['file']
            img = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), cv2.IMREAD_COLOR)

            # 使用全局模型进行推理
            results = async_predict(img)

            # 处理结果
            output = []
            for result in results:
                for box in result.boxes:
                    output.append({
                        'class': model.names[int(box.cls)],
                        'confidence': float(box.conf),
                        'box': box.xywhn.tolist()[0]
                    })

            return Response({'results': output}, status=status.HTTP_200_OK)

        except Exception as e:
            return Response({'error': str(e)}, status=status.HTTP_400_BAD_REQUEST)
